from copy import deepcopy
import random
from chiplet import QPU
from tool import remove_root_node_from_dag, extract_first_layer
import networkx as nx
from docplex.mp.model import Model
import math


class Highway:
    '''对应的Highway'''

    def __init__(self, serve_dag_id, generate_time, live_time, ghz_link, last_serve_time):
        self.serve_dag_id: int = serve_dag_id  # 服务的 CNOT 节点
        self.generate_time: int = generate_time  # 开始生成的时间
        self.live_time: int = live_time  # GHZ态的存续时间
        # 相关的ghz链路，tuple中是物理比特ID
        self.ghz_link: list[tuple[int, int]] = ghz_link
        self.last_serve_time: int = last_serve_time  # 上一次服务的完成时间


# ⭐️ 优化选路
def select_path_for_given_CNOTs(qpu: QPU, cnots: list[list], dag_ids: list[int]):
    '''针对给定的一批CNOT门，选择唯一的路径'''
    '''传出：一批路径、对应的dag节点ID'''
    '''示例：[[0,3,5,7,9], [1,2,5,6]...], [3,6...]'''

    candidate_paths = []

    max_candidate_num = 3
    for cnot in cnots:
        tmp = qpu.get_k_path_with_qubit_id_without_block(
            source_qubit=cnot[0], target_qubit=cnot[1], k=max_candidate_num)
        candidate_paths.append(tmp)

    cnot_length = len(cnots)
    conflict_graph = nx.Graph()

    for i in range(cnot_length):
        for j in range(max_candidate_num):
            conflict_graph.add_node((i, j))  # 使用 cnot索引 和 candidate索引 作为节点
            conflict_graph.nodes[(
                i, j)]['weight'] = qpu.circuit_dag.dag_nx.nodes[dag_ids[i]]["time_weight"]

    model = Model(name="select_max_weight")

    x_dim = [(i, j) for i in range(cnot_length)
             for j in range(max_candidate_num)]
    x = model.continuous_var_dict(x_dim, name="x", lb=0, ub=1)

    # 组内互斥
    for i in range(cnot_length):
        sum_expr = model.linear_expr()
        for j in range(max_candidate_num):
            sum_expr += x[i, j]
        model.add_constraint(sum_expr <= 1)

    # 构造冲突图
    for i in range(cnot_length):
        for j in range(max_candidate_num):
            for _i in range(i, cnot_length):
                for _j in range(max_candidate_num):
                    if (i != _i or j != _j):
                        path = candidate_paths[i][j]
                        _path = candidate_paths[_i][_j]
                        if (is_two_path_conflict(path, _path)):
                            conflict_graph.add_edge((i, j), (_i, _j))
                            model.add_constraint(x[i, j] + x[_i, _j] <= 1)

    objective_expr = model.linear_expr()
    for i in range(cnot_length):
        for j in range(max_candidate_num):
            objective_expr += x[i, j] * conflict_graph.nodes[(i, j)]['weight']

    model.maximize(objective_expr)

    selected_paths = []
    selected_dag_nodes = []

    sol = model.solve()
    if sol:
        while conflict_graph.number_of_nodes() > 0:
            max_ij_x = -1
            max_ij = (0, 0)
            for i in range(cnot_length):
                for j in range(max_candidate_num):
                    if (not conflict_graph.has_node((i, j))):
                        continue
                    tmp_multiple = x[i, j].solution_value * \
                        conflict_graph.nodes[(i, j)]['weight']
                    if tmp_multiple > max_ij_x:
                        max_ij_x = tmp_multiple
                        max_ij = (i, j)

            # 移除该节点以及与该节点冲突的所有节点
            neighbors = list(conflict_graph.neighbors(max_ij))
            conflict_graph.remove_node(max_ij)
            for __tmp in neighbors:
                conflict_graph.remove_node(__tmp)

            # 将最大的加入路径列表
            selected_paths.append(candidate_paths[max_ij[0]][max_ij[1]])
            selected_dag_nodes.append(dag_ids[max_ij[0]])

    return selected_paths, selected_dag_nodes


def shuttle_for_cnot(qpu: QPU, cnot_paths, dag_node_ids, final_time_dag_optimal: nx.DiGraph):
    '''⭐️ 优化调度'''
    '''给定：完整的QPU、CNOT路径数组、DAG节点ID数组、DAG图'''
    '''优化节点之间的调度，给DAG的节点打上时间标签'''
    '''返回：所需要的额外的Qubit数量'''
    path_num = len(cnot_paths)
    ancilla_qubit_num = 0
    for i in range(path_num):
        dag_node_id = dag_node_ids[i]
        cnot_path = cnot_paths[i]
        max_pre_time = 0
        for pre in final_time_dag_optimal.predecessors(dag_node_id):
            max_pre_time = max(
                max_pre_time, final_time_dag_optimal.nodes[pre]["process_done_time"])
        def ceil_int_x_linear(L):
            x = 0
            while 2**x + 2*x < L:
                x += 1
            return x
        L = len(cnot_path)
        eliminated_num = ceil_int_x_linear(L)
        cross_num = get_cross_computing_qubit_num(cnot_path, qpu)
        final_time_dag_optimal.nodes[dag_node_id]["process_done_time"] = max_pre_time + eliminated_num + cross_num
        ancilla_qubit_num += cross_num + L - 2 * eliminated_num
    return ancilla_qubit_num


def schedule_optimal(original_qpu: QPU, callbackSchedulePath=None):
    '''⭐️ 优化调度，计算最后电路完成时间'''
    qpu: QPU = deepcopy(original_qpu)  # 先拷贝一份

    qpu.estimate_critical_value()  # 预估权重
    # ⭐️ 在调度的过程中，量子比特的映射会发生变化
    final_time_dag_optimal = deepcopy(qpu.circuit_dag.dag_nx)

    tmp_dag_optimal = deepcopy(qpu.circuit_dag.dag_nx)
    remove_root_node_from_dag(tmp_dag_optimal)
    has_remaining_cnot = True

    remain_GHZ_highway: list[Highway] = []  # 保留残存的GHZ Highway

    # swap_num = 0
    # ghz_length = 0
    ancilla_num = 0

    while has_remaining_cnot:
        # 抽取 Front Layer
        cnots_in_window, tmp_dag_with_attrs, dag_node_id_in_window, _has_remaining_cnot = extract_first_layer(
            tmp_dag_optimal)

        # 判断DAG图中是否有剩余的CNOT门没处理
        has_remaining_cnot = _has_remaining_cnot
        if (not has_remaining_cnot):
            break
        
        # ⭐️ 选路，这些路都是 non_conflict 的，以及他们的节点ID
        cnot_paths, dag_node_ids = select_path_for_given_CNOTs(
            qpu, cnots_in_window, dag_node_id_in_window)

        # ⭐️ 调度路径
        ancilla_qubit_num = shuttle_for_cnot(
            qpu, cnot_paths, dag_node_ids, final_time_dag_optimal)

        ancilla_num += ancilla_qubit_num
        
        tmp_dag_optimal.remove_nodes_from(dag_node_ids)

    max_depth = 0
    # 找出 self.final_depth_dag 出度为0的顶点
    for node in final_time_dag_optimal.nodes():
        if (final_time_dag_optimal.nodes[node]["process_done_time"] > max_depth):
            max_depth = final_time_dag_optimal.nodes[node]["process_done_time"]

    return max_depth, ancilla_num


def get_cross_computing_qubit_num(path, qpu: QPU):
    '''获取穿过了的量子比特数'''
    count = 0
    for i in path:
        if (qpu.map_plan[i] != None):
            count += 1
    return count


def is_two_path_conflict(path1, path2) -> bool:
    # 使用set来判断是否冲突
    set_path1 = set(path1)
    for node in path2:
        if node in set_path1:
            return True
    return False
